"""
Author: Xinyuan Ye
Description: Heuristic-Based Sudoku Solver Visualization using MRV and Forward Checking.
             This version uses the Minimum Remaining Values (MRV) heuristic and 
             forward checking to optimize the backtracking algorithm, speeding up 
             the solving process. Valid numbers are shown in green, and invalid 
             numbers (leading to backtracking) are shown in red.

Usage:
1. Save the Sudoku puzzle as a text file (0 represents empty cells).
2. Run the program: python sudoku_visualization.py <sudoku_file.txt>
"""

import pygame
import time
import sys

# Define constants for the board
WIDTH, HEIGHT = 540, 540  # Window size
ROWS, COLS = 9, 9  # 9x9 Sudoku grid
CELL_SIZE = WIDTH // COLS  # Cell size
GRID_COLOR = (0, 0, 0)  # Black for grid lines
VALID_COLOR = (0, 255, 0)  # Green for valid number
INVALID_COLOR = (255, 0, 0)  # Red for invalid attempt
NUMBER_COLOR = (0, 0, 0)  # Default color for numbers

# Initialize pygame
pygame.init()

# Set the display size
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Sudoku Solver Visualization - CSP")

# Define the font
font = pygame.font.SysFont(None, 40)

# Read the Sudoku puzzle from a text file
def load_sudoku_from_file(filepath):
    with open(filepath, 'r') as f:
        puzzle = []
        for line in f:
            puzzle.append([int(num) for num in line.split()])
    return puzzle

# Draw the Sudoku board
def draw_board(board):
    # Draw the grid
    for i in range(ROWS + 1):
        thickness = 4 if i % 3 == 0 else 1  # Thicker lines for 3x3 boxes
        pygame.draw.line(screen, GRID_COLOR, (0, i * CELL_SIZE), (WIDTH, i * CELL_SIZE), thickness)
        pygame.draw.line(screen, GRID_COLOR, (i * CELL_SIZE, 0), (i * CELL_SIZE, HEIGHT), thickness)

    # Draw the numbers on the board
    for i in range(ROWS):
        for j in range(COLS):
            if board[i][j] != 0:
                text = font.render(str(board[i][j]), True, NUMBER_COLOR)
                screen.blit(text, (j * CELL_SIZE + 15, i * CELL_SIZE + 15))

# Improved: Check if placing num in board[row][col] is valid
def is_valid(board, num, row, col):
    # Check row
    for i in range(COLS):
        if board[row][i] == num:
            return False
    # Check column
    for i in range(ROWS):
        if board[i][col] == num:
            return False
    # Check 3x3 grid
    box_row = row // 3 * 3
    box_col = col // 3 * 3
    for i in range(3):
        for j in range(3):
            if board[box_row + i][box_col + j] == num:
                return False
    return True

# Heuristic: Get the next empty cell with the fewest legal values (MRV)
def get_next_cell(board):
    min_options = 10  # Maximum of 9 options, so we start with 10
    next_cell = None
    for row in range(ROWS):
        for col in range(COLS):
            if board[row][col] == 0:  # Find empty cell
                options = [num for num in range(1, 10) if is_valid(board, num, row, col)]
                if len(options) < min_options:
                    min_options = len(options)
                    next_cell = (row, col)
    return next_cell

# Solve the Sudoku board with backtracking and MRV heuristic
def solve_sudoku(board):
    cell = get_next_cell(board)
    if cell is None:  # No empty cell left, puzzle is solved
        return True

    row, col = cell
    for num in range(1, 10):
        if is_valid(board, num, row, col):
            board[row][col] = num
            draw_current_cell(row, col, num, VALID_COLOR)  # Show the number in green for valid
            pygame.display.update()
            time.sleep(0.01)  # Optional delay for speed control

            if solve_sudoku(board):  # Recursively try to solve
                return True
            board[row][col] = 0  # Reset the cell if solution fails
            draw_current_cell(row, col, num, INVALID_COLOR)  # Show the number in red for invalid
            pygame.display.update()
            time.sleep(0.1)  # Optional delay for speed control

    return False

# Highlight the current cell and show the number being tested
def draw_current_cell(row, col, num, color):
    # Draw cell background color
    pygame.draw.rect(screen, color, (col * CELL_SIZE, row * CELL_SIZE, CELL_SIZE, CELL_SIZE))
    # Draw the number being tested
    text = font.render(str(num), True, NUMBER_COLOR)
    screen.blit(text, (col * CELL_SIZE + 15, row * CELL_SIZE + 15))
    # Redraw grid lines to keep them visible
    pygame.draw.rect(screen, GRID_COLOR, (col * CELL_SIZE, row * CELL_SIZE, CELL_SIZE, CELL_SIZE), 1)

# Main visualization loop
def main():
    if len(sys.argv) < 2:
        print("Usage: python sudoku_visualization.py <sudoku_file.txt>")
        sys.exit(1)

    # Load the Sudoku puzzle from a text file
    filepath = sys.argv[1]
    sudoku_board = load_sudoku_from_file(filepath)

    running = True
    screen.fill((255, 255, 255))  # White background
    draw_board(sudoku_board)
    pygame.display.update()

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        solve_sudoku(sudoku_board)
        draw_board(sudoku_board)
        pygame.display.update()

    pygame.quit()

if __name__ == "__main__":
    main()